import matplotlib
import os
import pandas as pd

def get_data():
    demand_train_A = 'data/demand_train_A.csv'
    geo_topo = 'data/geo_topo.csv'
    inventory_info_A = 'data/inventory_info_A.csv'
    product_topo = 'data/product_topo.csv'
    weight_A = 'data/weight_A.csv'

    demand_train_A = pd.read_csv(demand_train_A)
    geo_topo = pd.read_csv(geo_topo)
    inventory_info_A = pd.read_csv(inventory_info_A)
    product_topo = pd.read_csv(product_topo)
    weight_A = pd.read_csv(weight_A)

    dfs = [demand_train_A,geo_topo,inventory_info_A,product_topo,weight_A]
    for df in dfs:
        if 'Unnamed: 0' in df.columns:
            df.drop(columns='Unnamed: 0',inplace=True)
        if 'ts' in df.columns:
            df = df.sort_values(by='ts')
    return dfs

def get_future(demand_pridict):
    monday_future_data = demand_pridict
    monday_future_data['yesterday_qty'] = monday_future_data.groupby('unit')['qty'].shift(1).fillna(method='ffill').reset_index().sort_index().set_index('index')
    monday_future_data['diff_1'] = monday_future_data['qty'] - monday_future_data['yesterday_qty']
    monday_future_data['qty'] = monday_future_data['diff_1']
    
    #原baseline求的是每周一对应十四天之后一周的需求增量，这里改成对应本周需求增量，决策（solver函数）时再考虑偏移量
    # monday_future_data['shift_14']=monday_future_data.groupby('unit')['qty'].shift(-14).fillna(0).reset_index().sort_index().set_index('index')
    # monday_future_data = monday_future_data[['unit','ts','shift_14']].rename(columns={'shift_14':'qty'}) 
    monday_future_data = monday_future_data[['unit','ts','qty']].reset_index().sort_index().set_index('index') 
    monday_future_data['dt'] = pd.to_datetime(monday_future_data['ts'])
    monday_future_data['weekofyear'] = monday_future_data['dt'].dt.weekofyear
    monday_future_data['year'] = monday_future_data['dt'].dt.year
    monday_future_data_week = monday_future_data.copy()

    monday_future_data_week = monday_future_data_week.groupby(['weekofyear','year','unit'],as_index=False).sum()
    monday_future_data_week['sum_qty'] = monday_future_data_week['qty']
    monday_future_data = pd.merge(monday_future_data_week,monday_future_data,on = ['weekofyear','year','unit'])

    monday_future_data['dayofweek'] = monday_future_data['dt'].dt.dayofweek
    monday_future_data = monday_future_data[monday_future_data['dayofweek']==0]
    monday_future_data = monday_future_data[['unit','ts','sum_qty']].rename(columns={'sum_qty':'qty'})
    return monday_future_data

def solver(inventory_info_A,monday_future_data):
    init_inventory = inventory_info_A.set_index(['unit'])['qty'].to_dict()
    def consume_init_inventory(arr,init_val):
        remain = init_val
        i = 0
        result = []
        while  i <len(arr)-2:
            if i == 0:
                arr[i] = max(0,arr[i+2]+arr[i+1]+arr[i]-remain)
                remain -= arr[i]
                remain = max(remain,0)
            elif i == 1:     
                arr[i] = max(0,arr[i+2]+arr[i+1]+arr[i]-remain-arr[i-1])   
                remain -= arr[i]
                remain = max(remain,0)
            else:
                arr[i] = max(0,arr[i+2]+arr[i+1]+arr[i]-remain-arr[i-1])   
                remain =remain - arr[i] +arr[i-2]#第三周开始还要加上两周前的补货量
                remain = max(remain,0)
            result.append(arr[i])
            i+=1
        result += [0,0]
        
        return result

    r = []
    for i,group in monday_future_data.groupby('unit'):

        unit = group['unit'].values[0]
        init_val = init_inventory[unit]

        group = group.sort_values(by='ts')
        qty_list = group['qty'].values
        print(qty_list) 
        print(init_val)
        qty_list = consume_init_inventory(qty_list,init_val)
        group['qty'] = qty_list
        r.append(group)
    submission = pd.concat(r)
    return submission

if __name__ == "__main__":
    data = get_data()
    future_data = pd.read_csv('data/demand_test_A.csv')
    monday_future_data = get_future(future_data)

    submission = solver(data[2],monday_future_data)
    submission.to_csv("submission.csv",index=False)
    

